# encoding:utf-8
import networkx as nx


class UTG(object):
    def __init__(self, apk_name):
        self.belong_apk = apk_name
        self.G = nx.DiGraph()

    def add_transition(self, from_state, to_state, events):
        self.G.add_edge(from_state, to_state, events=events)

    def output_utg(self):
        utg = []
        for state_transition in self.G.edges():
            from_state = state_transition[0]
            to_state = state_transition[1]
            events = self.G[from_state][to_state]
            utg_edge = {
                "from": from_state.time_name,
                "to": to_state.time_name,
                "events": events
            }
            utg.append(utg_edge)
        return utg